{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile MAS.py\n",
    "from __future__ import print_function, division\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "import matplotlib.pyplot as pltD\n",
    "import time\n",
    "import copy\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "sys.path.append('General_utils')\n",
    "sys.path.append('MAS_utils')\n",
    "from ImageFolderTrainVal import *\n",
    "\n",
    "from MAS_based_Training import *\n",
    "\n",
    "from test_network import *\n",
    "import pdb\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "\n",
    "def exp_lr_scheduler(optimizer, epoch, init_lr=0.0004, lr_decay_epoch=54):\n",
    "    \"\"\"Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs.\"\"\"\n",
    "   \n",
    "    lr = init_lr * (0.1**(epoch // lr_decay_epoch))\n",
    "\n",
    "    if epoch % lr_decay_epoch == 0:\n",
    "        print('LR is set to {}'.format(lr))\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "    return optimizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "def MAS(dataset_path,previous_task_model_path,exp_dir,data_dir,reg_sets,reg_lambda=1,norm='L2', num_epochs=100,lr=0.0008,batch_size=200,b1=False):\n",
    "    \"\"\"Call MAS on mainly a sequence of two object recognition tasks with a head for each \n",
    "    dataset_path=new dataset path\n",
    "    exp_dir=where to save the new model\n",
    "    previous_task_model_path: previous task in the sequence model path to start from \n",
    "    reg_sets,data_dir: sets of examples used to compute omega. Here the default is the training set of the last task\n",
    "    b1=to mimic online importance weight computation, batch size=1\n",
    "    reg_lambda= regulizer hyper parameter. In object recognition it was set to 1.\n",
    "    norm=the norm used to compute the gradient of the learned function\n",
    "    \"\"\"\n",
    "    dsets = torch.load(dataset_path)\n",
    "    dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=150,\n",
    "                                               shuffle=True, num_workers=4)\n",
    "                for x in ['train', 'val']}\n",
    "    dset_sizes = {x: len(dsets[x]) for x in ['train', 'val']}\n",
    "    dset_classes = dsets['train'].classes\n",
    "    model_ft = torch.load(previous_task_model_path)\n",
    "    use_gpu = torch.cuda.is_available()\n",
    "\n",
    "    #update omega value\n",
    "    if b1:\n",
    "        update_batch_size=1\n",
    "    else:\n",
    "        update_batch_size=batch_size\n",
    "    model_ft=update_weights_params(data_dir,reg_sets,model_ft,update_batch_size,norm)\n",
    "        \n",
    "    #set the lambda value for the MAS    \n",
    "    model_ft.reg_params['lambda']=reg_lambda\n",
    "    \n",
    "    #get the number of features in this network and add a new task head\n",
    "    last_layer_index=str(len(model_ft.classifier._modules)-1)\n",
    "\n",
    "    num_ftrs=model_ft.classifier._modules[last_layer_index].in_features \n",
    "    model_ft.classifier._modules[last_layer_index] = nn.Linear(num_ftrs, len(dset_classes))    \n",
    "   \n",
    "    #check the computed omega\n",
    "    sanitycheck(model_ft)\n",
    "    \n",
    "    #define the loss function\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "   \n",
    "    \n",
    "    if use_gpu:\n",
    "        model_ft = model_ft.cuda()\n",
    "    \n",
    " \n",
    "    #call the MAS optimizer\n",
    "    optimizer_ft =Weight_Regularized_SGD(model_ft.parameters(), lr, momentum=0.9)\n",
    "    \n",
    "    if not os.path.exists(exp_dir):\n",
    "        os.makedirs(exp_dir)\n",
    "    resume=os.path.join(exp_dir,'epoch.pth.tar')\n",
    "    #train the model\n",
    "    #this training functin passes the reg params to the optimizer to be used for penalizing changes on important params\n",
    "    model_ft =train_model(model_ft, criterion, optimizer_ft,exp_lr_scheduler, lr,dset_loaders,dset_sizes,use_gpu,num_epochs,exp_dir,resume)\n",
    "    \n",
    "    return model_ft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "def MAS_sequence(dataset_path,pevious_pathes,previous_task_model_path,exp_dir,data_dirs,reg_sets,reg_lambda=1,norm='L2', num_epochs=100,lr=0.0008,batch_size=200,weight_decay=1e-5,b1=False,after_freeze=1):\n",
    "    \"\"\"Call MAS on mainly a sequence of tasks with a head for each where between at each step it sees samples from all the previous tasks to approximate the importance weights \n",
    "    dataset_path=new dataset path\n",
    "    exp_dir=where to save the new model\n",
    "    previous_task_model_path: previous task in the sequence model path to start from \n",
    "    pevious_pathes:pathes of previous methods to use the previous heads in the importance weights computation. We assume that each task head is not changed in classification setup of different tasks\n",
    "    reg_sets,data_dirs: sets of examples used to compute omega. Here the default is the training set of the last task\n",
    "    b1=to mimic online importance weight computation, batch size=1\n",
    "    reg_lambda= regulizer hyper parameter. In object recognition it was set to 1.\n",
    "    norm=the norm used to compute the gradient of the learned function\n",
    "    \"\"\"\n",
    "    dsets = torch.load(dataset_path)\n",
    "    dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=150,\n",
    "                                               shuffle=True, num_workers=4)\n",
    "                for x in ['train', 'val']}\n",
    "    dset_sizes = {x: len(dsets[x]) for x in ['train', 'val']}\n",
    "    dset_classes = dsets['train'].classes\n",
    "\n",
    "    use_gpu = torch.cuda.is_available()\n",
    "   \n",
    "    model_ft = torch.load(previous_task_model_path)\n",
    "    \n",
    "    if b1:\n",
    "        update_batch_size=1\n",
    "    else:\n",
    "        update_batch_size=batch_size\n",
    "    model_ft=update_sequence_MAS_weights(data_dirs,reg_sets,pevious_pathes,model_ft,update_batch_size,norm)\n",
    "\n",
    "    model_ft.reg_params['lambda']=reg_lambda\n",
    "    #model_ft = torchvision.models.alexnet()\n",
    "    last_layer_index=str(len(model_ft.classifier._modules)-1)\n",
    "\n",
    "    num_ftrs=model_ft.classifier._modules[last_layer_index].in_features \n",
    "    model_ft.classifier._modules[last_layer_index] = nn.Linear(num_ftrs, len(dset_classes))  \n",
    "   \n",
    "    #check the values of omega\n",
    "    sanitycheck(model_ft)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    #update the objective based params\n",
    "    \n",
    "    if use_gpu:\n",
    "        model_ft = model_ft.cuda()\n",
    "    \n",
    "\n",
    "\n",
    "    #our optimizer\n",
    "    optimizer_ft =Weight_Regularized_SGD(model_ft.parameters(), lr, momentum=0.9,weight_decay=weight_decay)\n",
    "    #exp_dir='/esat/monkey/raljundi/pytorch/CUB11f_hebbian_finetuned'\n",
    "    if not os.path.exists(exp_dir):\n",
    "        os.makedirs(exp_dir)\n",
    "    resume=os.path.join(exp_dir,'epoch.pth.tar')\n",
    "    model_ft = train_model(model_ft, criterion, optimizer_ft,exp_lr_scheduler, lr,dset_loaders,dset_sizes,use_gpu,num_epochs,exp_dir,resume)\n",
    "    \n",
    "    return model_ft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "def MAS_Omega_Acuumelation(dataset_path,previous_task_model_path,exp_dir,data_dir,reg_sets,reg_lambda=1,norm='L2', num_epochs=100,lr=0.0008,batch_size=200,b1=True):\n",
    "    \"\"\"\n",
    "    In case of accumelating omega for the different tasks in the sequence, baisically to mimic the setup of standard methods where \n",
    "    the reguilizer is computed on the training set. Note that this doesn't consider our adaptation\n",
    "    dataset_path=new dataset path\n",
    "    exp_dir=where to save the new model\n",
    "    previous_task_model_path: previous task in the sequence model path to start from \n",
    "    reg_sets: sets of examples used to compute omega. Here the default is the training set of the last task\n",
    "    b1=to mimic online importance weight computation, batch size=1\n",
    "    reg_lambda= regulizer hyper parameter. In object recognition it was set to 1.\n",
    "    norm=the norm used to compute the gradient of the learned function\n",
    "    \"\"\"\n",
    "    dsets = torch.load(dataset_path)\n",
    "    dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=150,\n",
    "                                               shuffle=True, num_workers=4)\n",
    "                for x in ['train', 'val']}\n",
    "    dset_sizes = {x: len(dsets[x]) for x in ['train', 'val']}\n",
    "    dset_classes = dsets['train'].classes\n",
    "\n",
    "    use_gpu = torch.cuda.is_available()\n",
    "   \n",
    "    model_ft = torch.load(previous_task_model_path)\n",
    "        \n",
    "    if b1:\n",
    "        #compute the importance with batch size of 1, to mimic the online setting\n",
    "        update_batch_size=1\n",
    "    else:\n",
    "        update_batch_size=batch_size\n",
    "    #update the omega for the previous task, accumelate it over previous omegas    \n",
    "    model_ft=accumulate_MAS_weights(data_dir,reg_sets,model_ft,update_batch_size,norm)\n",
    "    #set the lambda for the MAS regularizer\n",
    "    model_ft.reg_params['lambda']=reg_lambda\n",
    "    \n",
    "    \n",
    "    #get the number of features in this network and add a new task head\n",
    "    last_layer_index=str(len(model_ft.classifier._modules)-1)\n",
    "\n",
    "    num_ftrs=model_ft.classifier._modules[last_layer_index].in_features \n",
    "    model_ft.classifier._modules[last_layer_index] = nn.Linear(num_ftrs, len(dset_classes))  \n",
    "   \n",
    "\n",
    "    \n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    #update the objective based params\n",
    "    \n",
    "    if use_gpu:\n",
    "        model_ft = model_ft.cuda()\n",
    "    \n",
    "\n",
    "\n",
    "    #call the MAS optimizer\n",
    "    optimizer_ft =Weight_Regularized_SGD(model_ft.parameters(), lr, momentum=0.9)\n",
    "\n",
    "    if not os.path.exists(exp_dir):\n",
    "        os.makedirs(exp_dir)\n",
    "\n",
    "    #if there is a checkpoint to be resumed, in case where the training has stopped before on a given task    \n",
    "    resume=os.path.join(exp_dir,'epoch.pth.tar')\n",
    "    \n",
    "    #train the model\n",
    "    #this training functin passes the reg params to the optimizer to be used for penalizing changes on important params\n",
    "    model_ft = train_model(model_ft, criterion, optimizer_ft,exp_lr_scheduler, lr,dset_loaders,dset_sizes,use_gpu,num_epochs,exp_dir,resume)\n",
    "    \n",
    "    return model_ft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "def update_weights_params(data_dir,reg_sets,model_ft,batch_size,norm='L2'):\n",
    "    \"\"\"update the importance weights based on the samples included in the reg_set. Assume starting from zero omega\n",
    "    \n",
    "       model_ft: the model trained on the previous task \n",
    "    \"\"\"\n",
    "    data_transform =  transforms.Compose([\n",
    "        transforms.Scale(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n",
    "        \n",
    "    \n",
    "\n",
    "    #prepare the dataset\n",
    "    dset_loaders=[]\n",
    "    for data_path in reg_sets:\n",
    "    \n",
    "        # if so then the reg_sets is a dataset by its own, this is the case for the mnist dataset\n",
    "        if data_dir is not None:\n",
    "            dset = ImageFolderTrainVal(data_dir, data_path, data_transform)\n",
    "        else:\n",
    "            dset=torch.load(data_path)\n",
    "            dset=dset['train']\n",
    "        dset_loader= torch.utils.data.DataLoader(dset, batch_size=batch_size,\n",
    "                                               shuffle=False, num_workers=4)\n",
    "        dset_loaders.append(dset_loader)\n",
    "\n",
    "    \n",
    "    use_gpu = torch.cuda.is_available()\n",
    "\n",
    "    #inialize the importance params,omega, to zero\n",
    "    reg_params=initialize_reg_params(model_ft)\n",
    "    model_ft.reg_params=reg_params\n",
    "    #define the importance weight optimizer. Actually it is only one step. It can be integrated at the end of the first task training\n",
    "    optimizer_ft = MAS_Omega_update(model_ft.parameters(), lr=0.0001, momentum=0.9)\n",
    "    \n",
    "    if norm=='L2':\n",
    "        print('********************MAS with L2 norm***************')\n",
    "        #compute the imporance params\n",
    "        model_ft = compute_importance_l2(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "    else:\n",
    "        if norm=='vector':\n",
    "            optimizer_ft=MAS_Omega_Vector_Grad_update(model_ft.parameters(), lr=0.0001, momentum=0.9)\n",
    "\n",
    "            model_ft = compute_importance_gradient_vector(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "\n",
    "        else:\n",
    "            model_ft = compute_importance(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "\n",
    "   \n",
    "    return model_ft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "#\n",
    "def update_sequence_MAS_weights(data_dirs,reg_sets,previous_models,model_ft,batch_size,norm='L2'):\n",
    "    \"\"\"updates a task in a sequence while computing omega from scratch each time on the previous tasks\n",
    "       previous_models: to use their heads for compute the importance params\n",
    "    \"\"\"\n",
    "    data_transform =  transforms.Compose([\n",
    "        transforms.Scale(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n",
    "    t=0    \n",
    "    last_layer_index=str(len(model_ft.classifier._modules)-1)\n",
    "    for model_path in previous_models:\n",
    "        pre_model=torch.load(model_path)\n",
    "        #get previous task head\n",
    "        model_ft.classifier._modules[last_layer_index] = pre_model.classifier._modules[last_layer_index]\n",
    "\n",
    "        # if data_dirs is None then the reg_sets is a dataset by its own, this is the case for the mnist dataset\n",
    "        if data_dirs is not None:\n",
    "            dset = ImageFolderTrainVal(data_dirs[t], reg_sets[t], data_transform)\n",
    "        else:\n",
    "            dset=torch.load(reg_sets[t])\n",
    "            dset=dset['train']\n",
    "        dset_loader= torch.utils.data.DataLoader(dset, batch_size=batch_size,\n",
    "                                               shuffle=False, num_workers=4)\n",
    "       \n",
    "    #=============================================================================\n",
    "\n",
    "        use_gpu = torch.cuda.is_available()\n",
    "\n",
    "       \n",
    "        if t==0:\n",
    "            #initialize to zero\n",
    "            reg_params=initialize_reg_params(model_ft)\n",
    "        else:\n",
    "            #store previous task param\n",
    "            reg_params=initialize_store_reg_params(model_ft)\n",
    "        model_ft.reg_params=reg_params\n",
    "        #define the importance weight optimizer. Actually it is only one step. It can be integrated at the end of the first task training\n",
    "        optimizer_ft = MAS_Omega_update(model_ft.parameters(), lr=0.0001, momentum=0.9)\n",
    "        #legacy code\n",
    "        dset_loaders=[dset_loader]\n",
    "        #compute param importance\n",
    "        if norm=='L2':\n",
    "            \n",
    "            print('********************objective with L2 norm***************')\n",
    "            model_ft = compute_importance_l2(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "        else:\n",
    "            model_ft = compute_importance(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "        if t>0:\n",
    "            reg_params=accumelate_reg_params(model_ft)\n",
    "        \n",
    "        model_ft.reg_params=reg_params\n",
    "        t=t+1\n",
    "    sanitycheck(model_ft)   \n",
    "    return model_ft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "def accumulate_MAS_weights(data_dir,reg_sets,model_ft,batch_size,norm='L2'):\n",
    "    \"\"\"accumelate the importance params: stores the previously computed omega, compute omega on the last previous task\n",
    "            and accumelate omega resulting on  importance params for all the previous tasks\n",
    "       reg_sets:either a list of files containing the samples used for computing the importance param like train or train and test\n",
    "                or pytorch dataset, then train set is used\n",
    "       data_dir:\n",
    "    \"\"\"\n",
    "    data_transform =  transforms.Compose([\n",
    "        transforms.Scale(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n",
    "        \n",
    "    \n",
    "\n",
    "    #prepare the dataset\n",
    "    dset_loaders=[]\n",
    "    for data_path in reg_sets:\n",
    "    \n",
    "        # if data_dir is not None then the reg_sets is a dataset by its own, this is the case for the mnist dataset\n",
    "        if data_dir is not None:\n",
    "            dset = ImageFolderTrainVal(data_dir, data_path, data_transform)\n",
    "        else:\n",
    "            dset=torch.load(data_path)\n",
    "            dset=dset['train']\n",
    "        dset_loader= torch.utils.data.DataLoader(dset, batch_size=batch_size,\n",
    "                                               shuffle=False, num_workers=4)\n",
    "        dset_loaders.append(dset_loader)\n",
    "    #=============================================================================\n",
    "    \n",
    "    use_gpu = torch.cuda.is_available()\n",
    "\n",
    "    #store the previous omega, set values to zero\n",
    "    reg_params=initialize_store_reg_params(model_ft)\n",
    "    model_ft.reg_params=reg_params\n",
    "    #define the importance weight optimizer. Actually it is only one step. It can be integrated at the end of the first task training\n",
    "    optimizer_ft = MAS_Omega_update(model_ft.parameters(), lr=0.0001, momentum=0.9)\n",
    "   \n",
    "    if norm=='L2':\n",
    "        print('********************objective with L2 norm***************')\n",
    "        model_ft =compute_importance_l2(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "    else:\n",
    "        model_ft =compute_importance(model_ft, optimizer_ft,exp_lr_scheduler, dset_loaders,use_gpu)\n",
    "    #accumelate the new importance params  with the prviously stored ones (previous omega)\n",
    "    reg_params=accumelate_reg_params(model_ft)\n",
    "    model_ft.reg_params=reg_params\n",
    "    sanitycheck(model_ft)   \n",
    "    return model_ft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appending to MAS.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile -a MAS.py\n",
    "\n",
    "\n",
    "def sanitycheck(model):\n",
    "    for name, param in model.named_parameters():\n",
    "           \n",
    "            print (name)\n",
    "            if param in model.reg_params:\n",
    "            \n",
    "                reg_param=model.reg_params.get(param)\n",
    "                omega=reg_param.get('omega')\n",
    "                \n",
    "                print('omega max is',omega.max())\n",
    "                print('omega min is',omega.min())\n",
    "                print('omega mean is',omega.mean())"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
